
import caffe
import numpy as np
import lmdb
import os

root_file_path = "/home/sunzy/workspace/pyhome/CaffeProject/CaffeProject/pulsar/deep_learning/temp/"
deploy = root_file_path + 'mnist/deploy.prototxt'
caffe_model = root_file_path + 'mnist/custom_net_iter_1000.caffemodel'

def test():
    net = caffe.Net(deploy, caffe_model, caffe.TEST)
    img_lmdb = lmdb.open(os.path.join(root_file_path, "mnist/mnist_test_lmdb"))
    txn = img_lmdb.begin()
    cursor = txn.cursor()
    datum = caffe.proto.caffe_pb2.Datum()
    num = 0
    all = 0
    p = 0
    r = 0
    for (idx, (key, value)) in enumerate(cursor):
        datum.ParseFromString(value)
        flat_x = np.fromstring(datum.data, dtype=np.uint8)
        # print(datum.channels)
        x = flat_x.reshape( datum.channels, datum.height, datum.width)
        y = datum.label
        num += 1




        # x2 = transformer.preprocess("data", x)
        # print(x.shape)
        # print(x2.shape)
        # print(x2)
        # print(net.blobs.keys())
        # if num == 1:
        #     break


        net.blobs['data'].data[...] = x
        out = net.forward()

        prob = out['Softmax1']
        order = prob.argmax()
        print(num, y, order)
        # print(num, order, prob)
        if y == order:
            all+=1


    print(num, all)




def main():
    test()

if __name__ == "__main__":
    main()